"""
Case Type   : DATAVEC模块
Case Name   : 创建ivfpq 余弦距离索引，索引选项pq_m参数覆盖，参数不同，查询准确性验证
Create At   : 2025/4/14
Owner       : @zhanghuan
Description :
    step1:创建3张表;
    step2:创建自定义插入函数;
    step3:表插入相同数据;
    step4:创建索引，不指定pq_m，默认值8;
    step5:创建索引，指定pq_m为40;
    step6:创建索引，指定pq_m为80;
    step7:不同pq_m值索引，查询比较;
    step8:不同pq_m值索引，insert500条数据;
    step9:不同pq_m值索引，查询比较
    step10:不同pq_m值索引，update数据;
    step11:不同pq_m值索引，查询比较;
    step12:不同pq_m值索引，delete数据;
    step13:不同pq_m值索引，查询比较;
    step14:清理测试数据;
Expect      :
    step1:创建3张表; expect:成功
    step2:创建自定义插入函数; expect:成功
    step3:表插入相同数据; expect:成功
    step4:创建索引，不指定pq_m，默认值8; expect:成功
    step5:创建索引，指定pq_m为40; expect:成功
    step6:创建索引，指定pq_m为80; expect:成功
    step7:不同pq_m值索引，查询比较; expect:pq_m值越大，查询结果越准确
    step8:不同pq_m值索引，insert500条数据; expect:成功
    step9:不同pq_m值索引，查询比较; expect:pq_m值越大，查询结果越准确
    step10:不同pq_m值索引，update数据; expect:成功
    step11:不同pq_m值索引，查询比较; expect:pq_m值越大，查询结果越准确
    step12:不同pq_m值索引，delete数据; expect:成功
    step13:不同pq_m值索引，查询比较; expect:pq_m值越大，查询结果越准确
    step14:清理测试数据; expect:成功
History     :
"""
import os
import random
import unittest
from decimal import Decimal

from testcase.utils.Common import Common
from testcase.utils.CommonSH import CommonSH
from testcase.utils.CommonSQL import CommonSQL
from testcase.utils.Constant import Constant
from testcase.utils.Logger import Logger
from yat.test import Node
from yat.test import macro

pri_dbuser = Node(node='PrimaryDbUser')
pq_lib_path = '/usr/local/sra_recall/lib/'
log = Logger()
log.info("----判断本地环境是否存在pq插件----")
check_cmd = f'if [ -d {pq_lib_path} ]; ' \
            f'then echo "True"; else echo "False"; fi'
log.info(check_cmd)
check_result = pri_dbuser.sh(check_cmd).result()
log.info(check_result)


@unittest.skipIf(check_result == 'False',
                 '----环境不具备PQ插件，跳过执行用例----')
class DataVecTestCase(unittest.TestCase):
    def setUp(self):
        self.sh = CommonSH('PrimaryDbUser')
        self.sta_1_sh = CommonSH('Standby1DbUser')
        self.sta_2_sh = CommonSH('Standby2DbUser')
        self.constant = Constant()
        self.com = Common()
        self.pq_env_name = 'DATAVEC_PQ_LIB_PATH'
        self.tb_name = 't_datavec_index_pq_case0088'
        self.id_name = 'id_datavec_index_pq_case0088'
        self.comsql = CommonSQL('PrimaryDbUser')
        self.fun_insert = 'f_datavec_index_pq_case0088'

    def test_main(self):
        step_txt = '----step1:创建3张表; expect:成功----'
        log.info(step_txt)
        create_tb = f"drop table if exists {self.tb_name}_1,{self.tb_name}_2," \
                    f"{self.tb_name}_3;" \
                    f"create table {self.tb_name}_1 " \
                    f"(id bigserial primary key, val vector(160));" \
                    f"create table {self.tb_name}_2 " \
                    f"(id bigserial primary key, val vector(160));" \
                    f"create table {self.tb_name}_3 " \
                    f"(id bigserial primary key, val vector(160));"
        result = self.sh.execut_db_sql(create_tb)
        log.info(result)
        self.assertIn(self.constant.TABLE_CREATE_SUCCESS, result,
                      '执行失败:' + step_txt)

        step_txt = '----step2:创建自定义插入函数; expect:成功----'
        log.info(step_txt)
        result = self.comsql.create_vector_insert_func(self.fun_insert)
        log.info(result)
        self.assertIn(self.fun_insert, result, '执行失败:' + step_txt)

        step_txt = '----step3:表插入相同数据; expect:成功----'
        log.info(step_txt)
        insert_sql = f"select {self.fun_insert}(160,5000,'{self.tb_name}_1'," \
                     f"'val','vector');"
        result = self.sh.execut_db_sql(insert_sql)
        log.info(result)
        self.assertIn('SUCCESS', result, '执行失败:' + step_txt)
        result = self.sh.execut_db_sql(insert_sql)
        log.info(result)
        self.assertIn('SUCCESS', result, '执行失败:' + step_txt)
        sql = f"insert into {self.tb_name}_2 select * from {self.tb_name}_1;"
        result = self.sh.execut_db_sql(sql)
        log.info(result)
        self.assertIn(self.constant.INSERT_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)
        sql = f"insert into {self.tb_name}_3 select * from {self.tb_name}_1;"
        result = self.sh.execut_db_sql(sql)
        log.info(result)
        self.assertIn(self.constant.INSERT_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)

        step_txt = '----step4:创建索引，不指定pq_m，默认值8; expect:成功----'
        log.info(step_txt)
        create_index = (f"set maintenance_work_mem='100MB';"
                        f"create index {self.id_name}_1 on {self.tb_name}_1 "
                        f"using ivfflat (val vector_cosine_ops) "
                        f"with(enable_pq=on);")
        result = self.sh.execut_db_sql(create_index)
        log.info(result)
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)

        step_txt = '----step5:创建索引，指定pq_m为40; expect:成功----'
        log.info(step_txt)
        create_index = (f"set maintenance_work_mem='100MB';"
                        f"create index {self.id_name}_2 on {self.tb_name}_2 "
                        f"using ivfflat (val vector_cosine_ops) "
                        f"with(enable_pq=on,pq_m=40);")
        result = self.sh.execut_db_sql(create_index)
        log.info(result)
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)

        step_txt = '----step6:创建索引，指定pq_m为80; expect:成功---'
        log.info(step_txt)
        create_index = (f"set maintenance_work_mem='100MB';"
                        f"create index {self.id_name}_3 on {self.tb_name}_3 "
                        f"using ivfflat (val vector_cosine_ops) "
                        f"with(enable_pq=on,pq_m=80);")
        result = self.sh.execut_db_sql(create_index)
        log.info(result)
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)

        step_txt = ('----step7:不同pq_m值索引，查询比较; '
                    'expect:pq_m值越大，查询结果越准确----')
        log.info(step_txt)
        result1 = []
        result2 = []
        result3 = []
        for i in range(1, 100):
            log.info(f'----第{i}次查询----')
            radom_vector = random.randint(1, 10000)
            sql = f"select id, val <-> (select val from {self.tb_name}_1 " \
                  f"where id ={radom_vector}) from {self.tb_name}_1 " \
                  f"order by val <-> (select val from {self.tb_name}_1 where" \
                  f" id ={radom_vector}) limit 10;"
            result1.append(self.comsql.vector_query_recall(
                sql, 'set ivfflat_probes=100;', is_detail=False))
            sql = f"select id, val <-> (select val from {self.tb_name}_2 " \
                  f"where id ={radom_vector}) from {self.tb_name}_2 " \
                  f"order by val <-> (select val from {self.tb_name}_2 where" \
                  f" id ={radom_vector}) limit 10;"
            result2.append(self.comsql.vector_query_recall(
                sql, 'set ivfflat_probes=100;', is_detail=False))
            sql = f"select id, val <-> (select val from {self.tb_name}_3 " \
                  f"where id ={radom_vector}) from {self.tb_name}_3 " \
                  f"order by val <-> (select val from {self.tb_name}_3 where" \
                  f" id ={radom_vector}) limit 10;"
            result3.append(self.comsql.vector_query_recall(
                sql, 'set ivfflat_probes=100;', is_detail=False))
        log.info(result1)
        total = sum(Decimal(str(x)) for x in result1)
        average1 = total / len(result1)
        log.info(average1)

        log.info(result2)
        total = sum(Decimal(str(x)) for x in result2)
        average2 = total / len(result2)
        log.info(average2)

        log.info(result3)
        total = sum(Decimal(str(x)) for x in result3)
        average3 = total / len(result3)
        log.info(average3)
        self.assertGreaterEqual(average3, average2, '执行失败:' + step_txt)
        self.assertGreaterEqual(average2, average1, '执行失败:' + step_txt)

        step_txt = '---step8:不同pq_m值索引，insert500条数据 expect:成功---'
        log.info(step_txt)
        insert_sql = f"select {self.fun_insert}(160,500,'{self.tb_name}_1'," \
                     f"'val','vector');"
        result = self.sh.execut_db_sql(insert_sql)
        log.info(result)
        self.assertIn('SUCCESS', result, '执行失败:' + step_txt)
        result = self.sh.execut_db_sql(insert_sql)
        log.info(result)
        self.assertIn('SUCCESS', result, '执行失败:' + step_txt)
        sql = (f"insert into {self.tb_name}_2 select * from {self.tb_name}_1 "
               f"where id>10000;")
        result = self.sh.execut_db_sql(sql)
        log.info(result)
        self.assertIn(self.constant.INSERT_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)
        sql = (f"insert into {self.tb_name}_3 select * from {self.tb_name}_1 "
               f"where id>10000;")
        result = self.sh.execut_db_sql(sql)
        log.info(result)
        self.assertIn(self.constant.INSERT_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)

        step_txt = ('----step9:不同pq_m值索引，查询比较; '
                    'expect:pq_m值越大，查询结果越准确----')
        log.info(step_txt)
        result1 = []
        result2 = []
        result3 = []
        for i in range(1, 100):
            log.info(f'----第{i}次查询----')
            radom_vector = random.randint(1, 10000)
            sql = f"select id, val <-> (select val from {self.tb_name}_1 " \
                  f"where id ={radom_vector}) from {self.tb_name}_1 " \
                  f"order by val <-> (select val from {self.tb_name}_1 where" \
                  f" id ={radom_vector}) limit 10;"
            result1.append(self.comsql.vector_query_recall(
                sql, 'set ivfflat_probes=100;', is_detail=False))
            sql = f"select id, val <-> (select val from {self.tb_name}_2 " \
                  f"where id ={radom_vector}) from {self.tb_name}_2 " \
                  f"order by val <-> (select val from {self.tb_name}_2 where" \
                  f" id ={radom_vector}) limit 10;"
            result2.append(self.comsql.vector_query_recall(
                sql, 'set ivfflat_probes=100;', is_detail=False))
            sql = f"select id, val <-> (select val from {self.tb_name}_3 " \
                  f"where id ={radom_vector}) from {self.tb_name}_3 " \
                  f"order by val <-> (select val from {self.tb_name}_3 where" \
                  f" id ={radom_vector}) limit 10;"
            result3.append(self.comsql.vector_query_recall(
                sql, 'set ivfflat_probes=100;', is_detail=False))
        log.info(result1)
        total = sum(Decimal(str(x)) for x in result1)
        average1 = total / len(result1)
        log.info(average1)

        log.info(result2)
        total = sum(Decimal(str(x)) for x in result2)
        average2 = total / len(result2)
        log.info(average2)

        log.info(result3)
        total = sum(Decimal(str(x)) for x in result3)
        average3 = total / len(result3)
        log.info(average3)
        self.assertGreaterEqual(average3, average2, '执行失败:' + step_txt)
        self.assertGreaterEqual(average2, average1, '执行失败:' + step_txt)

        step_txt = '----step10:不同pq_m值索引，update数据; expect:成功----'
        log.info(step_txt)
        sql = f"update {self.tb_name}_1 set val = val + " \
              f"array_fill(10,array[160])::vector where id>5000;"
        result = self.comsql.get_query_duration(sql)
        log.info(result)
        self.assertIn(self.constant.UPDATE_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)

        sql = f"update {self.tb_name}_2 set val = val + " \
              f"array_fill(10,array[160])::vector where id>5000;"
        result = self.comsql.get_query_duration(sql)
        log.info(result)
        self.assertIn(self.constant.UPDATE_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)
        sql = f"update {self.tb_name}_3 set val = val + " \
              f"array_fill(10,array[160])::vector where id>5000;"
        result = self.comsql.get_query_duration(sql)
        log.info(result)
        self.assertIn(self.constant.UPDATE_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)

        step_txt = '----step11:reindex; expect:成功----'
        log.info(step_txt)
        sql = f"reindex table {self.tb_name}_1;"
        result = self.sh.execut_db_sql(sql)
        log.info(result)
        self.assertIn('REINDEX', result, '执行失败:' + step_txt)
        sql = f"reindex table {self.tb_name}_2;"
        result = self.sh.execut_db_sql(sql)
        log.info(result)
        self.assertIn('REINDEX', result, '执行失败:' + step_txt)
        sql = f"reindex table {self.tb_name}_3;"
        result = self.sh.execut_db_sql(sql)
        log.info(result)
        self.assertIn('REINDEX', result, '执行失败:' + step_txt)

        step_txt = ('----step12:不同pq_m值索引，查询比较; '
                    'expect:pq_m值越大，查询结果越准确----')
        log.info(step_txt)
        result1 = []
        result2 = []
        result3 = []
        for i in range(1, 100):
            log.info(f'----第{i}次查询----')
            radom_vector = random.randint(1, 10000)
            sql = f"select id, val <-> (select val from {self.tb_name}_1 " \
                  f"where id ={radom_vector}) from {self.tb_name}_1 " \
                  f"order by val <-> (select val from {self.tb_name}_1 where" \
                  f" id ={radom_vector}) limit 10;"
            result1.append(self.comsql.vector_query_recall(
                sql, 'set ivfflat_probes=100;', is_detail=False))
            sql = f"select id, val <-> (select val from {self.tb_name}_2 " \
                  f"where id ={radom_vector}) from {self.tb_name}_2 " \
                  f"order by val <-> (select val from {self.tb_name}_2 where" \
                  f" id ={radom_vector}) limit 10;"
            result2.append(self.comsql.vector_query_recall(
                sql, 'set ivfflat_probes=100;', is_detail=False))
            sql = f"select id, val <-> (select val from {self.tb_name}_3 " \
                  f"where id ={radom_vector}) from {self.tb_name}_3 " \
                  f"order by val <-> (select val from {self.tb_name}_3 where" \
                  f" id ={radom_vector}) limit 10;"
            result3.append(self.comsql.vector_query_recall(
                sql, 'set ivfflat_probes=100;', is_detail=False))
        log.info(result1)
        total = sum(Decimal(str(x)) for x in result1)
        average1 = total / len(result1)
        log.info(average1)

        log.info(result2)
        total = sum(Decimal(str(x)) for x in result2)
        average2 = total / len(result2)
        log.info(average2)

        log.info(result3)
        total = sum(Decimal(str(x)) for x in result3)
        average3 = total / len(result3)
        log.info(average3)
        self.assertGreaterEqual(average3, average2, '执行失败:' + step_txt)
        self.assertGreaterEqual(average2, average1, '执行失败:' + step_txt)

        step_txt = '----step13:不同pq_m值索引，delete数据; expect:成功----'
        log.info(step_txt)
        sql = f"delete from {self.tb_name}_1 where id<5000;"
        result = self.comsql.get_query_duration(sql)
        log.info(result)
        self.assertIn(self.constant.DELETE_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)
        sql = f"delete from {self.tb_name}_2 where id<5000;"
        result = self.comsql.get_query_duration(sql)
        log.info(result)
        self.assertIn(self.constant.DELETE_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)
        sql = f"delete from {self.tb_name}_3 where id<5000;"
        result = self.comsql.get_query_duration(sql)
        log.info(result)
        self.assertIn(self.constant.DELETE_SUCCESS_MSG, result[1],
                      '执行失败:' + step_txt)

        step_txt = ('----step14:不同pq_m值索引，查询比较; '
                    'expect:pq_m值越大，查询结果越准确----')
        log.info(step_txt)
        result1 = []
        result2 = []
        result3 = []
        for i in range(1, 100):
            log.info(f'----第{i}次查询----')
            radom_vector = random.randint(5001, 10000)
            sql = f"select id, val <-> (select val from {self.tb_name}_1 " \
                  f"where id ={radom_vector}) from {self.tb_name}_1 " \
                  f"order by val <-> (select val from {self.tb_name}_1 where" \
                  f" id ={radom_vector}) limit 10;"
            result1.append(self.comsql.vector_query_recall(
                sql, 'set ivfflat_probes=100;', is_detail=False))
            sql = f"select id, val <-> (select val from {self.tb_name}_2 " \
                  f"where id ={radom_vector}) from {self.tb_name}_2 " \
                  f"order by val <-> (select val from {self.tb_name}_2 where" \
                  f" id ={radom_vector}) limit 10;"
            result2.append(self.comsql.vector_query_recall(
                sql, 'set ivfflat_probes=100;', is_detail=False))
            sql = f"select id, val <-> (select val from {self.tb_name}_3 " \
                  f"where id ={radom_vector}) from {self.tb_name}_3 " \
                  f"order by val <-> (select val from {self.tb_name}_3 where" \
                  f" id ={radom_vector}) limit 10;"
            result3.append(self.comsql.vector_query_recall(
                sql, 'set ivfflat_probes=100;', is_detail=False))
        log.info(result1)
        total = sum(Decimal(str(x)) for x in result1)
        average1 = total / len(result1)
        log.info(average1)

        log.info(result2)
        total = sum(Decimal(str(x)) for x in result2)
        average2 = total / len(result2)
        log.info(average2)

        log.info(result3)
        total = sum(Decimal(str(x)) for x in result3)
        average3 = total / len(result3)
        log.info(average3)
        self.assertGreaterEqual(average3, average2, '执行失败:' + step_txt)
        self.assertGreaterEqual(average2, average1, '执行失败:' + step_txt)

    def tearDown(self):
        step_txt = '----this is teardown----'
        log.info(step_txt)
        step1_txt = '----teardown:清理测试数据; expect:成功----'
        log.info(step1_txt)
        drop_cmd = f"drop table if exists {self.tb_name}_1,{self.tb_name}_2," \
                   f"{self.tb_name}_3;"
        drop_result1 = self.sh.execut_db_sql(drop_cmd)
        log.info(drop_result1)
        drop_cmd = f"drop function if exists {self.fun_insert};"
        drop_result2 = self.sh.execut_db_sql(drop_cmd)
        log.info(drop_result2)

        log.info(f'----{os.path.basename(__file__)}:end----')
        self.assertIn(self.constant.TABLE_DROP_SUCCESS, drop_result1,
                      '执行失败:' + step_txt + step1_txt)
        self.assertIn(self.constant.DROP_FUNCTION_SUCCESS_MSG, drop_result2,
                      '执行失败:' + step_txt + step1_txt)
